"""
/* Copyright 2018 The Enflame Tech Company. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
"""
# !/usr/bin/python
# coding=utf-8

import tensorflow as tf

slim = tf.contrib.slim

from enflame_models.enflame_alexnet import EnflameAlexnet
from enflame_models.alexnet import alexnet_v2, alexnet_v2_arg_scope
from enflame_models.resnet import resnet_arg_scope, resnet_v2_50, resnet_v2_14, resnet_v2_6, resnet_v2_101, \
    resnet_v2_152, resnet_v2_200
from enflame_models.vgg import vgg_arg_scope, vgg_11, vgg_16, vgg_19
from enflame_models.mnist import mnist_model
from enflame_models.resnet_model import resnet_cifar_model, resnet_imagenet_model
from enflame_models.inception_v1 import inception_v1, inception_v1_arg_scope
from enflame_models.inception_v2 import inception_v2, inception_v2_arg_scope
from enflame_models.inception_v3 import inception_v3, inception_v3_arg_scope
from enflame_models.inception_v4 import inception_v4, inception_v4_arg_scope
from enflame_models import mobilenet_v1
from enflame_models import mobilenet_v2
from enflame_models import mobilenet_v3
from utils.dtu_logger import LOGGER as logger


def build_model(x, params, is_training=True, weight_init=None, dropout_rate=1., reuse=tf.AUTO_REUSE):
    # Series of resnet networks have been configurable by depth and version
    # And not in use according to demand from QA that distinguishes model name in the name of logs
    end_points = None

    if params['model'] == 'enflame_alexnet':
        model = EnflameAlexnet(x, rate=dropout_rate, num_classes=params['num_class'], is_training=is_training)
        net = model.fc3
    elif params['model'] == 'alexnet' or params['model'] == 'std_alexnet':
        params['model'] = 'std_alexnet'
        with slim.arg_scope(alexnet_v2_arg_scope()):
            net, end_points = alexnet_v2(x, num_classes=params['num_class'], is_training=is_training,
                                         data_format=params['data_format'], use_resource=params['use_resource'])
    elif params['model'] == 'resnet':
        with slim.arg_scope(resnet_arg_scope(weight_init=weight_init)):
            net, end_points = resnet_v2_50(x, num_classes=params['num_class'], is_training=is_training,
                                           data_format=params['data_format'], use_resource=params['use_resource'],
                                           reuse=reuse)
    elif params['model'] == 'resnet6':
        with slim.arg_scope(resnet_arg_scope(weight_init=weight_init)):
            net, end_points = resnet_v2_6(x, num_classes=params['num_class'], is_training=is_training,
                                          data_format=params['data_format'], use_resource=params['use_resource'],
                                          reuse=reuse)
    elif params['model'] == 'resnet14':
        with slim.arg_scope(resnet_arg_scope(weight_init=weight_init)):
            net, end_points = resnet_v2_14(x, num_classes=params['num_class'], is_training=is_training,
                                           data_format=params['data_format'], use_resource=params['use_resource'],
                                           reuse=reuse)
    elif params['model'] == 'resnet101':
        with slim.arg_scope(resnet_arg_scope(weight_init=weight_init)):
            net, end_points = resnet_v2_101(x, num_classes=params['num_class'], is_training=is_training,
                                            data_format=params['data_format'], use_resource=params['use_resource'],
                                            reuse=reuse)
    elif params['model'] == 'resnet152':
        with slim.arg_scope(resnet_arg_scope(weight_init=weight_init)):
            net, end_points = resnet_v2_152(x, num_classes=params['num_class'], is_training=is_training,
                                            data_format=params['data_format'], use_resource=params['use_resource'],
                                            reuse=reuse)
    elif params['model'] == 'resnet200':
        with slim.arg_scope(resnet_arg_scope(weight_init=weight_init)):
            net, end_points = resnet_v2_200(x, num_classes=params['num_class'], is_training=is_training,
                                            data_format=params['data_format'], use_resource=params['use_resource'],
                                            reuse=reuse)
    elif params['model'] in ['googlenet', 'inception_v1']:
        with slim.arg_scope(inception_v1_arg_scope()):
            net, end_points = inception_v1(x, num_classes=params['num_class'], is_training=is_training,
                                           dropout_keep_prob=dropout_rate, reuse=reuse,
                                           data_format=params['data_format'], use_resource=params['use_resource'])
    elif params['model'] == 'inception_v2':
        with slim.arg_scope(inception_v2_arg_scope()):
            net, end_points = inception_v2(x, num_classes=params['num_class'], is_training=is_training,
                                           data_format=params['data_format'], use_resource=params['use_resource'],
                                           reuse=reuse, dropout_keep_prob=dropout_rate)
    elif params['model'] == 'inception_v3':
        with slim.arg_scope(inception_v3_arg_scope()):
            net, end_points = inception_v3(x, num_classes=params['num_class'], is_training=is_training,
                                           data_format=params['data_format'], use_resource=params['use_resource'],
                                           dropout_keep_prob=dropout_rate)
    elif params['model'] == 'inception_v4':
        with slim.arg_scope(inception_v4_arg_scope()):
            net, end_points = inception_v4(x, num_classes=params['num_class'], is_training=is_training,
                                           dropout_keep_prob=dropout_rate)
    elif params['model'] == 'mobilenet_v1':
        with tf.contrib.slim.arg_scope(mobilenet_v1.mobilenet_v1_arg_scope()):
            net, endpoints = mobilenet_v1.mobilenet_v1(x, num_classes=params['num_class'])
    elif params['model'] == 'mobilenet_v2':
        with tf.contrib.slim.arg_scope(mobilenet_v2.training_scope()):
            net, endpoints = mobilenet_v2.mobilenet(x, num_classes=params['num_class'])
    elif params['model'] == 'mobilenet_v3':
        with tf.contrib.slim.arg_scope(mobilenet_v3.training_scope()):
            net, endpoints = mobilenet_v3.mobilenet(x, num_classes=params['num_class'])
    elif params['model'] == 'vgg11':
        with slim.arg_scope(vgg_arg_scope()):
            net, end_points = vgg_11(x, num_classes=params['num_class'], is_training=is_training,
                                     dropout_keep_prob=dropout_rate, data_format=params['data_format'],
                                     use_resource=params['use_resource'])
    elif params['model'] in ['vgg', 'vgg16']:
        with slim.arg_scope(vgg_arg_scope()):
            net, end_points = vgg_16(x, num_classes=params['num_class'], is_training=is_training,
                                     dropout_keep_prob=dropout_rate, data_format=params['data_format'],
                                     use_resource=params['use_resource'])
    elif params['model'] == 'vgg19':
        with slim.arg_scope(vgg_arg_scope()):
            net, end_points = vgg_19(x, num_classes=params['num_class'], is_training=is_training,
                                     dropout_keep_prob=dropout_rate, data_format=params['data_format'],
                                     use_resource=params['use_resource'])
    elif params['model'] == 'resnet_cifar':
        if params['data_format'] == 'NHWC':
            data_format = 'channels_last'
        elif params['data_format'] == 'NCHW':
            data_format = 'channels_first'
        elif params['data_format'] == 'CHNW':
            data_format = 'channels_chnw'
        else:
            raise Exception("Data_format {} not support".format(params['data_format']))

        net = resnet_cifar_model(x, 50, dtype=params['dtype'], num_classes=params['num_class'],
                                 data_format=data_format, resnet_version=2, is_training=is_training)
    elif params['model'] == 'resnet50_v1.5':
        if params['data_format'] == 'NHWC':
            data_format = 'channels_last'
        elif params['data_format'] == 'NCHW':
            data_format = 'channels_first'
        elif params['data_format'] == 'CHNW':
            data_format = 'channels_chnw'
        else:
            raise Exception("Data_format {} not support".format(params['data_format']))

        net = resnet_imagenet_model(x, 50, dtype=params['dtype'], num_classes=params['num_class'],
                                    data_format=data_format, resnet_version=1.5, is_training=is_training)
    elif params['model'] == 'resnet50_v1':
        if params['data_format'] == 'NHWC':
            data_format = 'channels_last'
        elif params['data_format'] == 'NCHW':
            data_format = 'channels_first'
        elif params['data_format'] == 'CHNW':
            data_format = 'channels_chnw'
        else:
            raise Exception("Data_format {} not support".format(params['data_format']))

        net = resnet_imagenet_model(x, 50, dtype=params['dtype'], num_classes=params['num_class'],
                                    data_format=data_format, resnet_version=1, is_training=is_training)
    elif params['model'] == 'resnet50_v2':
        if params['data_format'] == 'NHWC':
            data_format = 'channels_last'
        elif params['data_format'] == 'NCHW':
            data_format = 'channels_first'
        elif params['data_format'] == 'CHNW':
            data_format = 'channels_chnw'
        else:
            raise Exception("Data_format {} not support".format(params['data_format']))

        net = resnet_imagenet_model(x, 50, dtype=params['dtype'], num_classes=params['num_class'],
                                    data_format=data_format, resnet_version=2, is_training=is_training)
    elif params['model'] == 'resnet18_v2':
        if params['data_format'] == 'NHWC':
            data_format = 'channels_last'
        elif params['data_format'] == 'NCHW':
            data_format = 'channels_first'
        elif params['data_format'] == 'CHNW':
            data_format = 'channels_chnw'
        else:
            raise Exception("Data_format {} not support".format(params['data_format']))

        net = resnet_imagenet_model(x, 18, dtype=params['dtype'], num_classes=params['num_class'],
                                    data_format=data_format, resnet_version=2, is_training=is_training)
    elif params['model'] == 'lenet':
        net = mnist_model(x, keep_prob=dropout_rate, use_resource=params['use_resource'])
    else:
        logger.warn("Unsupported model name {}".format(params['model']))
        raise ValueError(
            'enflame model test error: Unsupported model{}, model in '
            '(enflame_alexnet, alexnet, vgg/vgg19/vgg11/vgg16 , inception_v1/2/3/4, resnet, resnet6, resnet14, '
            'resnet_cifar, resnet50_v1.5, resnet50_v1, resnet50_v2, resnet18_v2, googlenet , resnet101, resnet152, resnet200, lenet'
            ')'.format(params['model']))
    return net, end_points
